﻿module FsMgProtocol
open System
open FsBson
open System.IO
open System.Threading

let mutable nextId = 0

[<Flags>]
type OpCode =
    | OP_REPLY          = 1     //Reply to a client request. responseTo is set
    | OP_MSG            = 1000  //generic msg command followed by a string
    | OP_UPDATE         = 2001  //update document
    | OP_INSERT         = 2002  //insert new document
    | RESERVED          = 2003  //formerly used for OP_GET_BY_OID
    | OP_QUERY          = 2004  //query a collection
    | OP_GET_MORE       = 2005  //Get more data from a query. See Cursors
    | OP_DELETE         = 2006  //Delete documents
    | OP_KILL_CURSORS   = 2007  //Tell database client is done with a cursor 

type MsgHeader =
    struct
        val mutable messageLength : int // total message size, including this
        val mutable requestID : int     // identifier for this message
        val mutable responseTo : int    // requestID from the original request
        val mutable opCode : OpCode     // request type - see table below
    end

[<Flags>]
type UpdateFlags =
| None = 0b00
| Upsert = 0b01         //If set, the database will insert the supplied object into the collection if no matching document is found.
| MultiUpdate = 0b10    //If set, the database will update all matching objects in the collection. Otherwise only updates first matching doc.

[<Flags>]
type InsertFlags =
| None = 0b00
| ContinueOnError = 0b01

[<Flags>]
type QueryFlags =
| None              = 0b00          //Reserved must be 0
| TailableCursor    = 0b10          //Tailable means cursor is not closed when the last data is retrieved. Rather, the cursor marks the final object's position. You can resume using the cursor later, from where it was located, if more data were received. Like any "latent cursor", the cursor may become invalid at some point (CursorNotFound) – for example if the final object it references were deleted.
| SlaveOk           = 0b100         //Allow query of replica slave. Normally these return an error except for namespace "local".
//| OplogReplay     = 0b1000        //Internal replication use only - driver should not set
| NoCursorTimeout   = 0b10000       //The server normally times out idle cursors after an inactivity period (10 minutes) to prevent excess memory use. Set this option to prevent that.
| AwaitData         = 0b100000      //Use with TailableCursor. If we are at the end of the data, block for a while rather than returning no data. After a timeout period, we do return as normal.
| Exhaust           = 0b1000000     //Stream the data down full blast in multiple "more" packages, on the assumption that the client will fully read all data queried. Faster when you are pulling a lot of data and know you want to pull it all down. Note: the client is not allowed to not read all the data unless it closes the connection.
| Partial           = 0b10000000    //Get partial results from a mongos if some shards are down (instead of throwing an error)

[<Flags>]
type DeleteFlags =
| None = 0b00
| SingleRemove = 0b01       //If set, the database will remove only the first matching document in the collection. Otherwise all matching documents will be removed.

[<Flags>]
type ReplyFlags =
| None              = 0b00
| CursorNotFound    = 0b01      //Set when getMore is called but the cursor id is not valid at the server. Returned with zero results.
| QueryFailure      = 0b10      //Set when query failed. Results consist of one document containing an "$err" field describing the failure.
| ShardConfigStale  = 0b100     //Drivers should ignore this. Only mongos will ever see this set, in which case, it needs to update config from the server.
| AwaitCapable      = 0b1000    //Set when the server supports the AwaitData Query option. If it doesn't, a client should sleep a little between getMore's of a Tailable cursor. Mongod version 1.6 supports AwaitData and thus always sets AwaitCapable.

type OpUpdateMsg = {Id:int; Collection:string; Flags:UpdateFlags; Selector:Bson; Update:Bson}
type OpInsertMsg = {Id:int; Collection:string; Flags:InsertFlags; Documents:Bson seq}
type OpQueryMsg = {Id:int; Collection:string; Flags:QueryFlags; Skip:int; Take:int; Query:Bson; Fields:Bson option}
type OpGetMoreMsg = {Id:int; Collection:string; Take:int; Cursor:Int64}
type OpDeleteMsg = {Id:int; Collection:string; Flags:DeleteFlags; Selector:Bson}
type OpKillCursorsMsg = {Id:int; Cursors:Int64 seq}
type OpReplyMsg = {Id:int; Flags:ReplyFlags; Cursor:Int64; Start:int; Count:int; Results:Bson seq}

type Command =
| Insert of OpInsertMsg
| Update of OpUpdateMsg
| Delete of OpDeleteMsg
| Query of OpQueryMsg
| GetMore of OpGetMoreMsg
| KillCursors of OpKillCursorsMsg

let byteSize cmd =
    let headerSize = 16 //4*4
    let zeroField = 4
    let flagsField = 4
    let stringEnd = 1
    let int32Field = 4
    let int64Field = 8
    match cmd with
    | Insert ins ->     headerSize + flagsField + (utf8Size ins.Collection) + stringEnd 
                        + (ins.Documents |> Seq.sumBy(fun d->d.ByteSize))
    | Update up ->      headerSize + zeroField + (utf8Size up.Collection) + stringEnd 
                        + flagsField + up.Selector.ByteSize + up.Update.ByteSize
    | Delete del ->     headerSize + zeroField + (utf8Size del.Collection) + stringEnd 
                        + flagsField + del.Selector.ByteSize
    | Query q ->        headerSize + flagsField + (utf8Size q.Collection) + stringEnd 
                        + int32Field + int32Field + q.Query.ByteSize
                        + (match q.Fields with Some f -> f.ByteSize | None -> 0)
    | GetMore gm ->     headerSize + zeroField + (utf8Size gm.Collection) + stringEnd 
                        + int32Field + int64Field
    | KillCursors kc -> headerSize + zeroField + int32Field + (kc.Cursors |> Seq.length)*int64Field


let nextCmdId() = Interlocked.Increment &nextId

let writeHdr (str:BinaryWriter) cmd  (reqId:int) (opCode:int) =
    let size = byteSize cmd
    str.Write(size)
    str.Write(reqId)
    str.Write(0)
    str.Write(opCode)

let serialize cmd (stream:Stream) =
    let dgSize=byteSize cmd
    let dgBuffer = Array.create dgSize 0uy
    let dgStream = new MemoryStream(dgBuffer)
    let str = new BinaryWriter(dgStream)
//    let str = new BinaryWriter(stream)
    
    let cmdId = Interlocked.Increment &nextId

    match cmd with
    | Insert ins ->
        writeHdr str cmd cmdId ((int)OpCode.OP_INSERT)
        str.Write((int)ins.Flags)
        writeCString str ins.Collection
        for b in ins.Documents do
             serializeBson b str
    | Update up ->
        writeHdr str cmd cmdId ((int)OpCode.OP_UPDATE)
        str.Write(0)
        writeCString str up.Collection
        str.Write((int)up.Flags)
        serializeBson up.Selector str
        serializeBson up.Update str   
    | Delete del -> 
        writeHdr str cmd cmdId ((int)OpCode.OP_DELETE)
        str.Write(0)
        writeCString str del.Collection
        str.Write((int)del.Flags)
        serializeBson del.Selector str
    | Query q -> 
        writeHdr str cmd cmdId ((int)OpCode.OP_QUERY)
        str.Write((int)q.Flags)
        writeCString str q.Collection
        str.Write(q.Skip)
        str.Write(q.Take)
        serializeBson q.Query str
        match q.Fields with Some f -> serializeBson f str | None -> ()
    | GetMore gm -> ()
    | KillCursors kc -> ()

    str.Flush()

    stream.Write(dgBuffer,0,dgBuffer.Length)
    stream.Flush()


let deserializeReply (str:Stream) =
    let br = new BinaryReader(str)
    
    //hdr
    let size = br.ReadInt32()
    let requestId = br.ReadInt32()
    let responseTo = br.ReadInt32()
    let opCode = br.ReadInt32()
    if opCode <> 1 then failwith "unexpected opCode for reply"

    //reply
    let respFlags = br.ReadInt32()
    let cursorId = br.ReadInt64()
    let startingFrom = br.ReadInt32()
    let numberReturned = br.ReadInt32()

    let fixedSize = 36 //4+4+4+4+4+8+4+4

    let respSize = size - fixedSize
    let docs =
        match respSize with
        | 0 -> Array.empty 
        | s ->
            let bytes = br.ReadBytes(respSize)
            use ms = new MemoryStream(bytes)
            let br = new BinaryReader(ms)
            (0) 
            |> Seq.unfold(fun (count) ->
                if count < numberReturned then
                   let doc,_ = (|Document|) br
                   Some (BDoc(doc),(count+1))
                else
                    None) 
            |> Seq.toArray

    {Id=requestId; 
     Flags=ReplyFlags.AwaitCapable;
     Cursor=cursorId;
     Start=startingFrom;
     Count=numberReturned;
     Results=docs}
    


    
